import os
import xml.etree.ElementTree as ET
import shutil

def convert_voc_to_yolo(voc_path, yolo_path):
    classes = ['blue', 'white', 'yellow', 'red', 'none', 'correct', 'person', 'orange', 'wrong']
  # 替换为您实际的类别列表

    annotations_dir = os.path.join(voc_path, 'Annotations')
    image_sets_dir = os.path.join(voc_path, 'ImageSets', 'Main')
    jpeg_images_dir = os.path.join(voc_path, 'JPEGImages')

    train_txt = os.path.join(image_sets_dir, 'train.txt')
    yolo_labels_dir = os.path.join(yolo_path, 'labels', 'train')
    yolo_images_dir = os.path.join(yolo_path, 'images', 'train')

    with open(train_txt, 'r') as f:
        train_files = [line.strip() for line in f.readlines()]

    for filename in train_files:
        xml_file = os.path.join(annotations_dir, f"{filename}.xml")
        image_file = os.path.join(jpeg_images_dir, f"{filename}.jpg")
        yolo_label_file = os.path.join(yolo_labels_dir, f"{filename}.txt")

        tree = ET.parse(xml_file)
        root = tree.getroot()

        image_width = int(root.find('size/width').text)
        image_height = int(root.find('size/height').text)

        with open(yolo_label_file, 'w+') as f:
            for obj in root.findall('object'):
                class_name = obj.find('name').text
                class_id = classes.index(class_name)

                bbox = obj.find('bndbox')
                xmin = float(bbox.find('xmin').text)
                ymin = float(bbox.find('ymin').text)
                xmax = float(bbox.find('xmax').text)
                ymax = float(bbox.find('ymax').text)

                x = (xmin + xmax) / (2 * image_width)
                y = (ymin + ymax) / (2 * image_height)
                width = (xmax - xmin) / image_width
                height = (ymax - ymin) / image_height

                f.write(f"{class_id} {x} {y} {width} {height}\n")

        shutil.copy2(image_file, yolo_images_dir)
        print(f"Converted {filename}")

# 用法示例
voc_path = '/mnt/e/datasets/VOC'
yolo_path = '/mnt/e/datasets/YOLO'
convert_voc_to_yolo(voc_path, yolo_path)
